import cv2
import numpy as np


# 霍夫变换实现检测图像中的指定条直线
def Hough_Line(edge, img):
    # 进行变换，从x,y到rho，theta，并统计在变换域中交点的个数
    def voting(edge):
        H, W = edge.shape
        drho = 1
        dtheta = 1
        # 先求最大半径rho
        rho_max = np.ceil(np.sqrt(H ** 2 + W ** 2)).astype(np.int16)
        # 初始化矩阵
        hough = np.zeros((rho_max, 180), dtype=np.int16)
        # 获取边界
        # ind[0] 是 符合条件的纵坐标，ind[1]是符合条件的横坐标
        ind = np.where(edge == 255)
        ## hough变换函数
        # zip函数返回元组
        for y, x in zip(ind[0], ind[1]):
            for theta in range(0, 180, dtheta):
                # 转换到极坐标下
                t = np.pi / 180 * theta  # 换为弧度
                rho = int(x * np.cos(t) + y * np.sin(t))
                # 在hough域下，把该点对应的曲线（散点图）做出，dtheta取1---此处画线可理解为出现一次该点+1
                hough[rho, theta] += 1

        out = hough.astype(np.uint8)  # 截断处理

        return out

    # 非极大抑制
    def non_maximum_suppression(hough):
        rho_max, _ = hough.shape
        ##非极大抑制
        for y in range(rho_max):
            for x in range(180):
                # 查找附近区域，避免由于轮廓本身的大小导致同一处被画多条线
                x1 = max(x - 4, 0)
                x2 = min(x + 4, 180)
                y1 = max(y - 4, 0)
                y2 = min(y + 4, rho_max - 1)
                if np.max(hough[y1:y2, x1:x2]) == hough[y, x] and hough[y, x] != 0:
                    pass
                    # hough[y,x] = 255
                else:
                    hough[y, x] = 0  # 非极大抑制，当附近出现相同交点数目时，置零，避免轮廓本身大小对于hough检测的影响

        return hough

    def inverse_hough(hough, img):
        H, W, _ = img.shape
        rho_max, _ = hough.shape

        out = img.copy()

        # get x, y index of hough table
        # np.ravel 将多维数组降为1维
        # argsort  将数组元素从小到大排序，返回索引
        # [::-1]   反序->从大到小
        # [:10]    前10个
        ind_x = np.argsort(hough.ravel())[::-1][:17]
        ind_y = ind_x.copy()
        thetas = ind_x % 180
        rhos = ind_y // 180

        # 反变换回到x，y平面
        for theta, rho in zip(thetas, rhos):
            # theta[radian] -> angle[degree]
            t = np.pi / 180. * theta  # 变为弧度
            # hough -> (x,y)
            for x in range(W):
                if np.sin(t) != 0:
                    y = - (np.cos(t) / np.sin(t)) * x + (rho) / np.sin(t)
                    y = int(y)
                    if y >= H or y < 0:
                        continue
                    out[y, x] = [0, 0, 255]  # 标为红色
            for y in range(H):
                if np.cos(t) != 0:
                    x = - (np.sin(t) / np.cos(t)) * y + (rho) / np.cos(t)
                    x = int(x)
                    if x >= W or x < 0:
                        continue
                    out[y, x] = [0, 0, 255]  # 标为红色

        out = out.astype(np.uint8)

        return out

    # voting
    hough = voting(edge)

    # non maximum suppression
    hough = non_maximum_suppression(hough)

    # inverse hough
    out = inverse_hough(hough, img)

    return out

def hough_tool(img):
    if len(img.shape) == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    else:
        gray = img.copy()
    edge = cv2.Canny(gray, 50, 150)
    out = Hough_Line(edge, img).astype(np.uint8)
    return out


if __name__ == '__main__':
    # Read image
    img = cv2.imread("R-C.jpg").astype(np.float32)
    img1 = cv2.imread("R-C.jpg")
    img1 = cv2.GaussianBlur(img1, (3, 3), 0)
    gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    edge = cv2.Canny(gray, 50, 150)
    # Hough
    out = Hough_Line(edge, img)
    out = out.astype(np.uint8)
    # Save result
    cv2.imwrite("out.jpg", out)
    cv2.imshow("result", out)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
